import torch
from torch import nn
import math

"""
点积缩放注意力机制
"""


class DotProductAttention(nn.Module):
    def __init__(self, dropout):
        super().__init__()

        self.dropout = nn.Dropout(dropout)

    def forward(self, queries, keys, values, valid_lens):
        # queries (n,seq_len,m)
        # 获取输入Q的最后一个维度，作为缩放因子
        dim = queries.shape[-1]
        # 点积缩放机制： Qk.T / sqrt(d)
        scores = queries @ keys.T / torch.sqrt(dim)  # (n,seq_len,seq_len)
        # 返回带掩码的softmax结果 与 values点积得到最终的结果
        return masked_softmax(scores, valid_lens) @ values


def masked_softmax(scores, valid_lens, value=float(-math.inf)):  # valid_lens 掩码序列 torch.tenor([1,2,3])
    # 根据广播机制生成scores的掩码
    def _sequence_mask(scores, valid_lens):
        max_len = scores.shape[1]
        # 广播机制
        mask = torch.arange((max_len), dtype=torch.float32)[None, :] > valid_lens[:, None]
        mask = mask.float()
        mask[mask == 1] = value
        return mask

    if valid_lens is None:
        # 如果没有掩码，则返回原本的softmax的值
        return torch.softmax(scores, dim=-1)
    else:
        # 如果valid_lens只传递了一个数值
        if valid_lens.dims() == 1:  # 比如 valid_lens=3 那就改为 [3,3,3,3...]
            valid_lens = torch.repeat_interleave(valid_lens, scores.shape[0], dim=0)
        elif valid_lens.dims() > 2:  # 当维度过高时变为一维长序列
            valid_lens = valid_lens.reshape(-1)

        atten_mask = _sequence_mask(scores, valid_lens)
        return torch.softmax(scores + atten_mask, dim=-1)
